import cv2
import numpy as np
from rembg import remove
from PIL import Image
import mediapipe as mp


class BackgroundReplacerWithCrop:
    def __init__(self):
        self.bg_colors = {
            'white': (255, 255, 255),
            'blue': (219, 142, 67),  # BGR
            'red': (60, 20, 220),
            'gray': (128, 128, 128)
        }
        self.face_detector = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.6)

    def _remove_background(self, image):
        pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        result = remove(pil_img)
        return cv2.cvtColor(np.array(result), cv2.COLOR_RGBA2BGRA)

    def _change_background(self, image, bg_color):
        alpha = image[:, :, 3] / 255.0
        alpha = cv2.merge([alpha, alpha, alpha])

        background = np.full_like(image[:, :, :3], bg_color, dtype=np.uint8)
        foreground = image[:, :, :3].astype(float)
        background = background.astype(float)

        blended = cv2.multiply(alpha, foreground) + cv2.multiply(1.0 - alpha, background)
        return blended.astype(np.uint8)

    def _crop_face_proportionally(self, image, target_ratio=(4, 5), additional_top_margin_mm=3):
        """
        1. 使用去除背景后图像的最高点作为头顶
        2. 头顶与照片上边缘距离增加 2mm
        3. 头部高度占照片总高度的 2/3
        4. 双眼位于照片高度的 1/3 ~ 1/2 之间，目标为 0.4H
        """
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = self.face_detector.process(rgb)

        if not results.detections:
            raise RuntimeError("未检测到人脸")

        h, w = image.shape[:2]
        box = results.detections[0].location_data.relative_bounding_box
        x, y, bw, bh = box.xmin, box.ymin, box.width, box.height

        # 人脸框坐标（实际像素值）
        fx1 = int(x * w)
        fy1 = int(y * h)
        fx2 = int((x + bw) * w)
        fy2 = int((y + bh) * h)
        face_height = fy2 - fy1
        face_center_x = (fx1 + fx2) // 2

        # 设定头部占图像高度的 2/3，则裁剪高度为 face_height / (2/3)
        crop_h = int(face_height / (1 / 3))
        crop_w = int(crop_h * (target_ratio[0] / target_ratio[1]))

        # 假设双眼在脸框高度的 40%，即眼睛位置
        eye_y = fy1 + int(0.4 * face_height)

        # 目标眼睛位置应出现在最终图像的 0.4×H 高度处
        crop_top = eye_y - int(0.4 * crop_h)
        crop_bottom = crop_top + crop_h
        crop_left = face_center_x - crop_w // 2
        crop_right = crop_left + crop_w

        # 在头顶位置增加 2mm（单位转换为像素）
        additional_top_margin_pixels = int(additional_top_margin_mm / 25.4 * 300)  # 假设300DPI
        crop_top = max(crop_top - additional_top_margin_pixels, 0)

        # 边界限制，确保不会超出原图边界
        crop_left = max(crop_left, 0)
        crop_right = min(crop_right, w)
        crop_bottom = min(crop_bottom, h)

        # 裁剪图像
        crop_img = image[crop_top:crop_bottom, crop_left:crop_right]

        # 最后强制保证裁剪图为目标比例
        final_h, final_w = crop_img.shape[:2]
        current_ratio = final_w / final_h
        target_ratio_val = target_ratio[0] / target_ratio[1]

        if current_ratio > target_ratio_val:
            # 裁剪左右
            new_w = int(final_h * target_ratio_val)
            start_x = (final_w - new_w) // 2
            crop_img = crop_img[:, start_x:start_x + new_w]
        elif current_ratio < target_ratio_val:
            # 裁剪上下
            new_h = int(final_w / target_ratio_val)
            start_y = (final_h - new_h) // 2
            crop_img = crop_img[start_y:start_y + new_h, :]

        return crop_img

    def process(self, input_path, output_path, bg_color_name='white'):
        if bg_color_name not in self.bg_colors:
            raise ValueError(f"无效背景颜色，可选：{list(self.bg_colors.keys())}")

        img = cv2.imread(input_path)
        if img is None:
            raise FileNotFoundError(f"找不到图片：{input_path}")

        try:
            no_bg = self._remove_background(img)
            replaced_bg = self._change_background(no_bg, self.bg_colors[bg_color_name])
            cropped = self._crop_face_proportionally(replaced_bg)
            cv2.imwrite(output_path, cropped)
            print(f"✔ 已保存：{output_path}")
        except Exception as e:
            raise RuntimeError(f"❌ 处理失败：{str(e)}")


if __name__ == "__main__":
    processor = BackgroundReplacerWithCrop()

    processor.process("input.jpg", "person.jpg", bg_color_name="white")
    processor.process("input.jpg", "id_blue_crop.jpg", bg_color_name="blue")
    processor.process("input.jpg", "id_red_crop.jpg", bg_color_name="red")
